// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/vector/vector_maker.h>

namespace kumo::pollux {
    // static
    std::shared_ptr<const RowType> VectorMaker::row_type(
        std::vector<std::shared_ptr<const Type> > &&types) {
        std::vector<std::string> names;
        for (int32_t i = 0; i < types.size(); ++i) {
            names.push_back(fmt::format("c{}", i));
        }
        return ROW(std::move(names), std::move(types));
    }

    RowVectorPtr VectorMaker::row_vector(const std::vector<VectorPtr> &children) {
        std::vector<std::string> names;
        for (int32_t i = 0; i < children.size(); ++i) {
            names.push_back(fmt::format("c{}", i));
        }

        return row_vector(std::move(names), children);
    }

    RowVectorPtr VectorMaker::row_vector(
        std::vector<std::string> childNames,
        const std::vector<VectorPtr> &children) {
        std::vector<std::shared_ptr<const Type> > childTypes;
        childTypes.resize(children.size());
        for (int i = 0; i < children.size(); i++) {
            childTypes[i] = children[i]->type();
        }
        auto row_type = ROW(std::move(childNames), std::move(childTypes));
        const size_t vectorSize = children.empty() ? 0 : children.front()->size();

        return std::make_shared<RowVector>(
            pool_, row_type, BufferPtr(nullptr), vectorSize, children);
    }

    RowVectorPtr VectorMaker::row_vector(
        const std::shared_ptr<const RowType> &row_type,
        vector_size_t size) {
        return std::make_shared<RowVector>(
            pool_, row_type, BufferPtr(nullptr), size, std::vector<VectorPtr>{});
    }

    vector_size_t VectorMaker::create_offsets_and_sizes(
        vector_size_t size,
        std::function<vector_size_t(vector_size_t /* row */)> sizeAt,
        std::function<bool(vector_size_t /*row */)> is_null_at,
        BufferPtr *nulls,
        BufferPtr *offsets,
        BufferPtr *sizes) {
        *offsets = AlignedBuffer::allocate<vector_size_t>(size, pool_);
        *sizes = AlignedBuffer::allocate<vector_size_t>(size, pool_);

        auto rawOffsets = (*offsets)->asMutable<vector_size_t>();
        auto rawSizes = (*sizes)->asMutable<vector_size_t>();

        uint64_t *rawNulls;
        if (is_null_at) {
            *nulls = AlignedBuffer::allocate<bool>(size, pool_);
            rawNulls = (*nulls)->asMutable<uint64_t>();
            memset(rawNulls, bits::kNotNullByte, (*nulls)->capacity());
        }

        vector_size_t offset = 0;
        for (vector_size_t i = 0; i < size; ++i) {
            if (is_null_at && is_null_at(i)) {
                bits::set_null(rawNulls, i, true);
                rawOffsets[i] = 0;
                rawSizes[i] = 0;
                continue;
            }
            auto s = sizeAt(i);
            rawOffsets[i] = offset;
            rawSizes[i] = s;
            offset += s;
        }
        return offset;
    }

    ArrayVectorPtr VectorMaker::all_null_array_vector(
        vector_size_t size,
        const std::shared_ptr<const Type> &elementType) {
        BufferPtr nulls = allocate_nulls(size, pool_, bits::kNull);
        BufferPtr offsets = allocateOffsets(size, pool_);
        BufferPtr sizes = allocateSizes(size, pool_);

        return std::make_shared<ArrayVector>(
            pool_, ARRAY(elementType), nulls, size, offsets, sizes, nullptr);
    }

    MapVectorPtr VectorMaker::all_null_map_vector(
        vector_size_t size,
        const std::shared_ptr<const Type> &keyType,
        const std::shared_ptr<const Type> &valueType) {
        BufferPtr nulls = allocate_nulls(size, pool_, bits::kNull);
        BufferPtr offsets = allocateOffsets(size, pool_);
        BufferPtr sizes = allocateSizes(size, pool_);

        return std::make_shared<MapVector>(
            pool_,
            MAP(keyType, valueType),
            nulls,
            size,
            offsets,
            sizes,
            nullptr,
            nullptr);
    }

    namespace {
        template<TypeKind kind>
        void set_not_null_value(
            const VectorPtr &vector,
            vector_size_t index,
            variant value) {
            using T = typename TypeTraits<kind>::NativeType;

            if constexpr (std::is_same_v<T, StringView>) {
                vector->as_flat_vector<T>()->set(
                    index, StringView(value.value<const char *>()));
            } else {
                vector->as_flat_vector<T>()->set(index, value.value<T>());
            }
        }
    } // namespace

    ArrayVectorPtr VectorMaker::array_of_row_vector(
        const RowTypePtr &row_type,
        const std::vector<std::vector<variant> > &data) {
        vector_size_t size = data.size();
        BufferPtr offsets = AlignedBuffer::allocate<vector_size_t>(size, pool_);
        BufferPtr sizes = AlignedBuffer::allocate<vector_size_t>(size, pool_);

        auto rawOffsets = offsets->asMutable<vector_size_t>();
        auto rawSizes = sizes->asMutable<vector_size_t>();

        // Count number of elements.
        vector_size_t numElements = 0;
        for (const auto &array: data) {
            numElements += array.size();
        }

        // Create the underlying vector.
        auto row_vector = std::dynamic_pointer_cast<RowVector>(
            BaseVector::create(row_type, numElements, pool_));
        for (auto &field: row_vector->children()) {
            field->resize(numElements);
        }

        vector_size_t currentIdx = 0;
        for (const auto &arrayValue: data) {
            *rawSizes++ = arrayValue.size();
            *rawOffsets++ = currentIdx;

            for (auto arrayElement: arrayValue) {
                if (arrayElement.isNull()) {
                    row_vector->set_null(currentIdx, true);
                } else {
                    for (auto i = 0; i < row_type->size(); i++) {
                        const auto &value = arrayElement.row()[i];
                        if (value.isNull()) {
                            row_vector->childAt(i)->set_null(currentIdx, true);
                        } else {
                            POLLUX_DYNAMIC_SCALAR_TYPE_DISPATCH(
                                set_not_null_value,
                                row_type->childAt(i)->kind(),
                                row_vector->childAt(i),
                                currentIdx,
                                value);
                        }
                    }
                }
                currentIdx++;
            }
        }

        return std::make_shared<ArrayVector>(
            pool_,
            ARRAY(row_type),
            nullptr,
            size,
            offsets,
            sizes,
            std::move(row_vector),
            0);
    }

    ArrayVectorPtr VectorMaker::array_vector(
        const std::vector<vector_size_t> &offsets,
        const VectorPtr &elements,
        const std::vector<vector_size_t> &nulls) {
        const auto size = offsets.size();
        BufferPtr offsetsBuffer = allocateOffsets(size, pool_);
        BufferPtr sizesBuffer = allocateSizes(size, pool_);
        BufferPtr nullsBuffer = nullptr;
        auto rawOffsets = offsetsBuffer->asMutable<vector_size_t>();
        auto rawSizes = sizesBuffer->asMutable<vector_size_t>();

        for (int i = 0; i < size - 1; i++) {
            rawSizes[i] = offsets[i + 1] - offsets[i];
        }
        rawSizes[size - 1] = elements->size() - offsets.back();

        memcpy(rawOffsets, offsets.data(), size * sizeof(vector_size_t));

        if (!nulls.empty()) {
            nullsBuffer = AlignedBuffer::allocate<bool>(size, pool_, bits::kNotNull);
            auto rawNulls = nullsBuffer->asMutable<uint64_t>();

            for (int i = 0; i < nulls.size(); i++) {
                POLLUX_CHECK_EQ(rawSizes[nulls[i]], 0);
                bits::set_null(rawNulls, nulls[i]);
            }
        }

        return std::make_shared<ArrayVector>(
            pool_,
            ARRAY(elements->type()),
            nullsBuffer,
            size,
            offsetsBuffer,
            sizesBuffer,
            elements);
    }

    MapVectorPtr VectorMaker::map_vector(
        const std::vector<vector_size_t> &offsets,
        const VectorPtr &keys,
        const VectorPtr &values,
        const std::vector<vector_size_t> &nulls) {
        POLLUX_CHECK_EQ(keys->size(), values->size());

        const auto size = offsets.size();
        BufferPtr offsetsBuffer = allocateOffsets(size, pool_);
        BufferPtr sizesBuffer = allocateSizes(size, pool_);
        BufferPtr nullsBuffer = nullptr;
        auto rawOffsets = offsetsBuffer->asMutable<vector_size_t>();
        auto rawSizes = sizesBuffer->asMutable<vector_size_t>();

        for (int i = 0; i < size - 1; i++) {
            rawSizes[i] = offsets[i + 1] - offsets[i];
        }
        rawSizes[size - 1] = keys->size() - offsets.back();

        memcpy(rawOffsets, offsets.data(), size * sizeof(vector_size_t));

        if (!nulls.empty()) {
            nullsBuffer = AlignedBuffer::allocate<bool>(size, pool_, bits::kNotNull);
            auto rawNulls = nullsBuffer->asMutable<uint64_t>();

            for (int i = 0; i < nulls.size(); i++) {
                POLLUX_CHECK_EQ(rawSizes[nulls[i]], 0);
                bits::set_null(rawNulls, nulls[i]);
            }
        }

        return std::make_shared<MapVector>(
            pool_,
            MAP(keys->type(), values->type()),
            nullsBuffer,
            size,
            offsetsBuffer,
            sizesBuffer,
            keys,
            values);
    }
} // namespace kumo::pollux
